package xla

import (
	"github.com/gomlx/gomlx/backends"
	"github.com/gomlx/gopjrt/dtypes"
)

// CPUCapabilities supported by XLA CPU backends.
//
// This is the base value, and can be copied and specialized by specific PJRT that may not
// support everything.
var CPUCapabilities = backends.Capabilities{
	Operations: map[backends.OpType]bool{
		backends.OpTypeParameter:             true,
		backends.OpTypeConstant:              true,
		backends.OpTypeIdentity:              true,
		backends.OpTypeReduceWindow:          true,
		backends.OpTypeRNGBitGenerator:       true,
		backends.OpTypeBatchNormForInference: true,
		backends.OpTypeBatchNormForTraining:  true,
		backends.OpTypeBatchNormGradient:     true,
		backends.OpTypeBitCount:              true,

		backends.OpTypeAbs:                      true,
		backends.OpTypeAdd:                      true,
		backends.OpTypeArgMinMax:                true,
		backends.OpTypeBitcast:                  true,
		backends.OpTypeBitwiseAnd:               true,
		backends.OpTypeBitwiseNot:               true,
		backends.OpTypeBitwiseOr:                true,
		backends.OpTypeBitwiseXor:               true,
		backends.OpTypeBroadcast:                true,
		backends.OpTypeBroadcastInDim:           true,
		backends.OpTypeCeil:                     true,
		backends.OpTypeClz:                      true,
		backends.OpTypeComplex:                  true,
		backends.OpTypeConcatenate:              true,
		backends.OpTypeConj:                     true,
		backends.OpTypeConvGeneral:              true,
		backends.OpTypeConvertDType:             true,
		backends.OpTypeCos:                      true,
		backends.OpTypeDiv:                      true,
		backends.OpTypeDot:                      true,
		backends.OpTypeDotGeneral:               true,
		backends.OpTypeDynamicSlice:             true,
		backends.OpTypeDynamicUpdateSlice:       true,
		backends.OpTypeEqual:                    true,
		backends.OpTypeEqualTotalOrder:          true,
		backends.OpTypeErf:                      true,
		backends.OpTypeExp:                      true,
		backends.OpTypeExpm1:                    true,
		backends.OpTypeFFT:                      true,
		backends.OpTypeFloor:                    true,
		backends.OpTypeGather:                   true,
		backends.OpTypeGreaterOrEqual:           true,
		backends.OpTypeGreaterOrEqualTotalOrder: true,
		backends.OpTypeGreaterThan:              true,
		backends.OpTypeGreaterThanTotalOrder:    true,
		backends.OpTypeImag:                     true,
		backends.OpTypeIota:                     true,
		backends.OpTypeIsFinite:                 true,
		backends.OpTypeLessOrEqual:              true,
		backends.OpTypeLessOrEqualTotalOrder:    true,
		backends.OpTypeLessThan:                 true,
		backends.OpTypeLessThanTotalOrder:       true,
		backends.OpTypeLog:                      true,
		backends.OpTypeLog1p:                    true,
		backends.OpTypeLogicalAnd:               true,
		backends.OpTypeLogicalNot:               true,
		backends.OpTypeLogicalOr:                true,
		backends.OpTypeLogicalXor:               true,
		backends.OpTypeLogistic:                 true,
		backends.OpTypeMax:                      true,
		backends.OpTypeMin:                      true,
		backends.OpTypeMul:                      true,
		backends.OpTypeNeg:                      true,
		backends.OpTypeNotEqual:                 true,
		backends.OpTypeNotEqualTotalOrder:       true,
		backends.OpTypePad:                      true,
		backends.OpTypePow:                      true,
		backends.OpTypeReal:                     true,
		backends.OpTypeReduceBitwiseAnd:         true,
		backends.OpTypeReduceBitwiseOr:          true,
		backends.OpTypeReduceBitwiseXor:         true,
		backends.OpTypeReduceLogicalAnd:         true,
		backends.OpTypeReduceLogicalOr:          true,
		backends.OpTypeReduceLogicalXor:         true,
		backends.OpTypeReduceMax:                true,
		backends.OpTypeReduceMin:                true,
		backends.OpTypeReduceProduct:            true,
		backends.OpTypeReduceSum:                true,
		backends.OpTypeRem:                      true,
		backends.OpTypeReshape:                  true,
		backends.OpTypeReverse:                  true,
		backends.OpTypeRound:                    true,
		backends.OpTypeRsqrt:                    true,
		backends.OpTypeScatterMax:               true,
		backends.OpTypeScatterMin:               true,
		backends.OpTypeScatterSum:               true,
		backends.OpTypeSelectAndScatterMax:      true,
		backends.OpTypeSelectAndScatterMin:      true,
		backends.OpTypeSelectAndScatterSum:      true,
		backends.OpTypeShiftLeft:                true,
		backends.OpTypeShiftRightArithmetic:     true,
		backends.OpTypeShiftRightLogical:        true,
		backends.OpTypeSign:                     true,
		backends.OpTypeSin:                      true,
		backends.OpTypeSlice:                    true,
		backends.OpTypeSqrt:                     true,
		backends.OpTypeSub:                      true,
		backends.OpTypeTanh:                     true,
		backends.OpTypeTranspose:                true,
		backends.OpTypeWhere:                    true,
	},

	DTypes: map[dtypes.DType]bool{
		dtypes.Bool:       true,
		dtypes.Int8:       true,
		dtypes.Int16:      true,
		dtypes.Int32:      true,
		dtypes.Int64:      true,
		dtypes.Uint8:      true,
		dtypes.Uint16:     true,
		dtypes.Uint32:     true,
		dtypes.Uint64:     true,
		dtypes.Float16:    true,
		dtypes.Float32:    true,
		dtypes.Float64:    true,
		dtypes.BFloat16:   true,
		dtypes.Complex64:  true,
		dtypes.Complex128: true,
	},
}
